#include "ball_tracker.h"
#include "utils.h"  // 添加这个头文件
#include <iostream>
#include <iomanip>
#include <chrono>

BallTracker::BallTracker() 
    : tracknet_(nullptr), frame_count_(0), input_width_(0), input_height_(0), input_channels_(0),
      video_width_(0), video_height_(0), ratio_x_(1.0f), ratio_y_(1.0f),
      d_y_plane_(nullptr), d_uv_plane_(nullptr), d_rgb_buffer_(nullptr), h_rgb_buffer_(nullptr) {
    
    current_position_ = {0, 0, 0.0f, 0};
}

BallTracker::~BallTracker() {
    if (tracknet_) {
        tracknet_->destroy();
        delete tracknet_;
    }
    
    cleanupGPUBuffers();

    if (video_writer_.isOpened()) {
        video_writer_.release();
    }
    
    if (csv_file_.is_open()) {
        csv_file_.close();
    }
}

bool BallTracker::initGPUBuffers(int width, int height) {
    video_width_ = width;
    video_height_ = height;
    
    // 计算缓冲区大小
    size_t y_size = width * height;
    size_t uv_size = (width * height) / 2; // NV12格式
    
    std::cout << "Allocating GPU buffers for resolution: " << width << "x" << height << std::endl;
    std::cout << "Y plane size: " << y_size << " bytes" << std::endl;
    std::cout << "UV plane size: " << uv_size << " bytes" << std::endl;
    
    // 清理现有缓冲区
    cleanupGPUBuffers();
    
    // 分配GPU内存
    cudaError_t err;
    
    // 分配Y平面
    err = cudaMalloc((void**)&d_y_plane_, y_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate d_y_plane: " << cudaGetErrorString(err) << std::endl;
        return false;
    }
    std::cout << "Allocated Y plane on GPU: " << y_size << " bytes" << std::endl;

    // 分配UV平面
    err = cudaMalloc((void**)&d_uv_plane_, uv_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate d_uv_plane: " << cudaGetErrorString(err) << std::endl;
        cudaFree(d_y_plane_);
        d_y_plane_ = nullptr;
        return false;
    }
    std::cout << "Allocated UV plane on GPU: " << uv_size << " bytes" << std::endl;
    
    // 分配RGB缓冲区（3通道）
    size_t rgb_size = width * height * 3 * sizeof(float);
    err = cudaMalloc((void**)&d_rgb_buffer_, rgb_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate d_rgb_buffer: " << cudaGetErrorString(err) << std::endl;
        cleanupGPUBuffers();
        return false;
    }
    std::cout << "Allocated RGB buffer on GPU: " << rgb_size << " bytes" << std::endl;

    // 分配pinned memory用于主机-设备传输
    err = cudaMallocHost((void**)&h_rgb_buffer_, rgb_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate h_rgb_buffer (pinned memory): " << cudaGetErrorString(err) << std::endl;
        cleanupGPUBuffers();
        return false;
    }
    std::cout << "Allocated pinned memory for RGB: " << rgb_size << " bytes" << std::endl;
    
    // 初始化GPU内存
    err = cudaMemset(d_y_plane_, 0, y_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to initialize d_y_plane: " << cudaGetErrorString(err) << std::endl;
    }
    
    err = cudaMemset(d_uv_plane_, 128, uv_size); // UV平面通常初始化为128
    if (err != cudaSuccess) {
        std::cerr << "Failed to initialize d_uv_plane: " << cudaGetErrorString(err) << std::endl;
    }
    
    err = cudaMemset(d_rgb_buffer_, 0, rgb_size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to initialize d_rgb_buffer: " << cudaGetErrorString(err) << std::endl;
    }
    
    memset(h_rgb_buffer_, 0, rgb_size);
    
    std::cout << "GPU buffers allocated and initialized successfully" << std::endl;
    return true;
} 

void BallTracker::cleanupGPUBuffers() {
    if (d_y_plane_) {
        cudaFree(d_y_plane_);
        d_y_plane_ = nullptr;
    }
    if (d_uv_plane_) {
        cudaFree(d_uv_plane_);
        d_uv_plane_ = nullptr;
    }
    if (d_rgb_buffer_) {
        cudaFree(d_rgb_buffer_);
        d_rgb_buffer_ = nullptr;
    }
    if (h_rgb_buffer_) {
        cudaFreeHost(h_rgb_buffer_);
        h_rgb_buffer_ = nullptr;
    }
}

bool BallTracker::init(const std::string& engine_path) {
    tracknet_ = new TrackNetV3_TRT();
    
    if (!tracknet_->init(engine_path)) {
        std::cerr << "Failed to initialize TrackNetV3" << std::endl;
        delete tracknet_;
        tracknet_ = nullptr;
        return false;
    }
    
    // 获取模型输入尺寸
    tracknet_->getInputSize(input_width_, input_height_, input_channels_);
    std::cout << "TrackNetV3 input size: " << input_width_ << "x" << input_height_ 
              << "x" << input_channels_ << std::endl;
    
    return true;
}

bool BallTracker::processFrameGPU(VideoProcessor& video_processor) {
    if (!tracknet_) {
        std::cerr << "TrackNetV3 not initialized" << std::endl;
        return false;
    }
    
    if (!d_y_plane_ || !d_uv_plane_ || !d_rgb_buffer_ || !h_rgb_buffer_) {
        std::cerr << "GPU buffers not properly initialized" << std::endl;
        return false;
    }
    
    int stride_y = 0, stride_uv = 0;
    
    //std::cout << "Reading frame from GPU decoder..." << std::endl;

    // 从解码器读取YUV数据
    if (!video_processor.readFrameGPU(d_y_plane_, d_uv_plane_, stride_y, stride_uv)) {
        std::cerr << "Failed to read frame from GPU decoder" << std::endl;
        return false;
    }
    
    //std::cout << "Frame read successfully, stride_y: " << stride_y << ", stride_uv: " << stride_uv << std::endl;
    
    // 检测是否为 JPEG 范围的格式
    const VideoInfo& video_info = video_processor.getVideoInfo();
    bool is_jpeg_range = (video_info.pixel_format == AV_PIX_FMT_YUVJ420P);
    
    //std::cout << "Pixel format: " << av_get_pix_fmt_name(video_info.pixel_format) 
    //          << ", is_jpeg_range: " << is_jpeg_range << std::endl;
    
    //std::cout << "Converting YUV to RGB using CUDA..." << std::endl;
    
    // 使用CUDA核函数将YUV转换为RGB
    yuv2rgb(d_y_plane_, d_uv_plane_, stride_y, stride_uv, d_rgb_buffer_, video_width_, video_height_, is_jpeg_range);
    
    // 检查CUDA错误
    cudaError_t cuda_err = cudaGetLastError();
    if (cuda_err != cudaSuccess) {
        std::cerr << "CUDA error in yuv2rgb: " << cudaGetErrorString(cuda_err) << std::endl;
        return false;
    }
    
    // 等待GPU完成
    cuda_err = cudaDeviceSynchronize();
    if (cuda_err != cudaSuccess) {
        std::cerr << "CUDA synchronize error: " << cudaGetErrorString(cuda_err) << std::endl;
        return false;
    }
    
    //std::cout << "YUV to RGB conversion completed" << std::endl;
    
    // 将RGB数据从GPU复制到CPU
    size_t rgb_size = video_width_ * video_height_ * 3 * sizeof(float);
    //std::cout << "Copying RGB data from GPU to CPU, size: " << rgb_size << " bytes" << std::endl;
    
    cuda_err = cudaMemcpy(h_rgb_buffer_, d_rgb_buffer_, rgb_size, cudaMemcpyDeviceToHost);
    if (cuda_err != cudaSuccess) {
        std::cerr << "Failed to copy RGB data from device to host: " << cudaGetErrorString(cuda_err) << std::endl;
        return false;
    }
    
    //std::cout << "RGB data copied successfully" << std::endl;
    
    // 调试：检查前几个像素值
    //std::cout << "First few RGB values: ";
    //for (int i = 0; i < 12; i++) {
    //    std::cout << h_rgb_buffer_[i] << " ";
    //}
    //std::cout << std::endl;
    
    // 创建RGB帧 - 使用交错格式
    cv::Mat rgb_float_frame(video_height_, video_width_, CV_32FC3, h_rgb_buffer_);
    
    if (rgb_float_frame.empty()) {
        std::cerr << "Failed to create RGB float frame from buffer" << std::endl;
        return false;
    }
    
    //std::cout << "RGB float frame: " << rgb_float_frame.cols << "x" << rgb_float_frame.rows 
    //          << ", channels: " << rgb_float_frame.channels() << ", type: " << rgb_float_frame.type() << std::endl;
    
    // 转换为8位BGR
    cv::Mat bgr_frame;
    cv::Mat bgr_float_frame;
    
    // 首先从RGB转换为BGR
    cv::cvtColor(rgb_float_frame, bgr_float_frame, cv::COLOR_RGB2BGR);
    
    // 然后转换为8位
    bgr_float_frame.convertTo(bgr_frame, CV_8UC3, 255.0);
    
    if (bgr_frame.empty()) {
        std::cerr << "Failed to create BGR frame" << std::endl;
        return false;
    }
    
    //std::cout << "BGR frame: " << bgr_frame.cols << "x" << bgr_frame.rows 
    //          << ", channels: " << bgr_frame.channels() << std::endl;
    
    // 调试：保存第一帧用于检查
    if (frame_count_ == 0) {
        cv::imwrite("debug_first_frame.jpg", bgr_frame);
        std::cout << "Saved first frame for debugging: debug_first_frame.jpg" << std::endl;
    }
    
    // 调整尺寸以适应模型输入
    cv::Mat resized_frame;
    cv::resize(bgr_frame, resized_frame, cv::Size(input_width_, input_height_));
    
    // 添加到帧缓冲区
    frame_buffer_.push_back(resized_frame.clone());
    
    // 保持缓冲区大小为3帧
    if (frame_buffer_.size() > 3) {
        frame_buffer_.erase(frame_buffer_.begin());
    }
    
    bool processing_done = false;
    // 当有足够的帧时进行处理
    if (frame_buffer_.size() == 3) {
        //std::cout << "Processing frame buffer with TrackNetV3..." << std::endl;
        
        cv::Mat heatmap;
        auto start_time = std::chrono::high_resolution_clock::now();
        
        if (tracknet_->process(frame_buffer_, heatmap)) {
            auto end_time = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
            
            int ball_x, ball_y;
            if (tracknet_->getBallPosition(heatmap, ball_x, ball_y, 0.3f)) {
                ratio_x_ = static_cast<float>(video_width_) / input_width_;
                ratio_y_ = static_cast<float>(video_height_) / input_height_;
                
                current_position_.x = static_cast<int>(ball_x * ratio_x_);
                current_position_.y = static_cast<int>(ball_y * ratio_y_);
                current_position_.confidence = 1.0f;
                current_position_.frame_id = frame_count_;
                
                if (frame_count_ % 100 == 0) {
                    std::cout << "Frame " << frame_count_ << ": Ball detected at (" 
                              << current_position_.x << ", " << current_position_.y 
                              << ") - Processing time: " << duration.count() << "ms" << std::endl;
                }
            } else {
                current_position_.x = 0;
                current_position_.y = 0;
                current_position_.confidence = 0.0f;
                current_position_.frame_id = frame_count_;
                
                if (frame_count_ % 100 == 0) {
                    std::cout << "Frame " << frame_count_ << ": No ball detected" 
                              << " - Processing time: " << duration.count() << "ms" << std::endl;
                }
            }
            
            if (current_position_.confidence == 0.0){
                std::cout << "Clear track_history for no detected frame: " << frame_count_ << std::endl;
                track_history_.clear();
            } else{
                track_history_.push_back(current_position_);
            }
            processing_done = true;
            
        } else {
            std::cerr << "Failed to process frame with TrackNetV3" << std::endl;
            // 即使处理失败，也要记录帧
            current_position_.x = 0;
            current_position_.y = 0;
            current_position_.confidence = 0.0f;
            current_position_.frame_id = frame_count_;
            processing_done = true;
            
            return false;
        }
    } else {
        // 帧缓冲区不足3帧，记录空数据
        current_position_.x = 0;
        current_position_.y = 0;
        current_position_.confidence = 0.0f;
        current_position_.frame_id = frame_count_;
        
        if (frame_count_ % 100 == 0) {
            std::cout << "Frame " << frame_count_ << ": Buffer filling (" 
                      << frame_buffer_.size() << "/3 frames)" << std::endl;
        }
    }

    // 写入CSV文件 - 确保每帧都写入
    if (csv_file_.is_open()) {
        csv_file_ << frame_count_ << "," << current_position_.x << "," 
                 << current_position_.y << "," << current_position_.confidence << "\n";
        csv_file_.flush(); // 确保立即写入
        
        if (frame_count_ % 100 == 0) {
            std::cout << "Frame " << frame_count_ << " data written to CSV" << std::endl;
        }
    } else {
        std::cerr << "CSV file not open at frame " << frame_count_ << std::endl;
    }
    
    // 写入视频 - 使用原始的BGR帧
    if (video_writer_.isOpened()) {
        drawTrackingResult(bgr_frame);
        video_writer_.write(bgr_frame);
        
        if (frame_count_ % 100 == 0) {
            std::cout << "Frame " << frame_count_ << " written to video" << std::endl;
        }
    } else {
        std::cerr << "Video writer not open at frame " << frame_count_ << std::endl;
    }
    
    frame_count_++;
    return true;
} 

bool BallTracker::initVideoWriter(const std::string& video_path, int width, int height, int fps) {
    if (video_writer_.isOpened()) {
        video_writer_.release();
    }
    
    // 尝试不同的编码器
    std::vector<int> codecs = {
        cv::VideoWriter::fourcc('m', 'p', '4', 'v'),
        cv::VideoWriter::fourcc('a', 'v', 'c', '1'),
        cv::VideoWriter::fourcc('X', '2', '6', '4'),
        cv::VideoWriter::fourcc('H', '2', '6', '4'),
        cv::VideoWriter::fourcc('M', 'J', 'P', 'G')
    };
    
    for (int codec : codecs) {
        if (video_writer_.open(video_path, codec, fps, cv::Size(width, height))) {
            std::cout << "Video writer initialized with codec: " << codec 
                      << ", resolution: " << width << "x" << height 
                      << ", fps: " << fps << std::endl;
            return true;
        }
    }
    
    std::cerr << "Failed to initialize video writer with any codec" << std::endl;
    return false;
}

bool BallTracker::processFrameSoftware(VideoProcessor& video_processor) {
    if (!tracknet_) {
        std::cerr << "TrackNetV3 not initialized" << std::endl;
        return false;
    }
    
    cv::Mat frame;
    bool read_success = video_processor.readFrame(frame);
    if (!read_success) {
        std::cerr << "No frame read in video processor" << std::endl;
        return false;
    }
    
    if (frame.empty()) {
        std::cerr << "Read empty frame from video" << std::endl;
        return true; // 继续处理下一帧，而不是终止
    }
    
    //std::cout << "Read frame " << frame_count_ << " - size: " << frame.cols << "x" << frame.rows << std::endl;
    // 直接使用软件解码的帧进行处理
    bool process_successed =  processFrame(frame);
    return process_successed;
}

bool BallTracker::processFrame(const cv::Mat& frame) {
    if (!tracknet_) {
        std::cerr << "TrackNetV3 not initialized" << std::endl;
        return false;
    }
    
    // 调整帧尺寸以适应模型输入
    cv::Mat resized_frame;
    cv::resize(frame, resized_frame, cv::Size(input_width_, input_height_));
    
    // 添加到帧缓冲区
    frame_buffer_.push_back(resized_frame.clone());
    
    // 保持缓冲区大小为3帧
    if (frame_buffer_.size() > 3) {
        frame_buffer_.erase(frame_buffer_.begin());
    }
    
    std::cout << "Frame buffer size " << frame_buffer_.size() << std::endl;
    // 当有足够的帧时进行处理
    if (frame_buffer_.size() == 3) {
        cv::Mat heatmap;
        auto start_time = std::chrono::high_resolution_clock::now();
        
        if (tracknet_->process(frame_buffer_, heatmap)) {
            auto end_time = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
            
            int ball_x, ball_y;
            if (tracknet_->getBallPosition(heatmap, ball_x, ball_y, 0.3f)) {
                // 计算缩放比例
                ratio_x_ = static_cast<float>(frame.cols) / input_width_;
                ratio_y_ = static_cast<float>(frame.rows) / input_height_;
                
                current_position_.x = static_cast<int>(ball_x * ratio_x_);
                current_position_.y = static_cast<int>(ball_y * ratio_y_);
                current_position_.confidence = 1.0f;
                current_position_.frame_id = frame_count_;
                
                std::cout << "Frame " << frame_count_ << ": Ball detected at (" 
                          << current_position_.x << ", " << current_position_.y 
                          << ") - Inference time: " << duration.count() << "ms" << std::endl;
            } else {
                current_position_.x = 0;
                current_position_.y = 0;
                current_position_.confidence = 0.0f;
                current_position_.frame_id = frame_count_;
                
                std::cout << "Frame " << frame_count_ << ": No ball detected" 
                          << " - Inference time: " << duration.count() << "ms" << std::endl;
            }
            
            // 添加到轨迹历史
            if (current_position_.confidence == 0.0f){
                assert(false);
                track_history_.clear();
            } else {
                track_history_.push_back(current_position_);
            }
            
            // 限制轨迹历史大小
            if (track_history_.size() > 30) {
                assert(false);
                track_history_.erase(track_history_.begin());
            }
            
            // 写入CSV
            if (csv_file_.is_open()) {
                csv_file_ << frame_count_ << "," << current_position_.x << "," 
                         << current_position_.y << "," << current_position_.confidence << "\n";
            }
            
            frame_count_++;
            return true;
        } else {
            std::cerr << "Failed to process frame with TrackNetV3" << std::endl;
            return false;
        }
    } else {
         // 帧缓冲区不足3帧，返回true表示成功接收帧但还未处理
        std::cout << "Frame " << frame_count_ << " added to buffer (buffer size: " << frame_buffer_.size() << ")" << std::endl;
        
        // 创建临时位置信息
        current_position_.x = 0;
        current_position_.y = 0;
        current_position_.confidence = 0.0f;
        current_position_.frame_id = frame_count_;
        
        // 写入CSV（即使没有检测结果）
        if (csv_file_.is_open()) {
            csv_file_ << frame_count_ << ",0,0,0.0\n";
            csv_file_.flush();
        }
        
        frame_count_++;
        return true;
    }
}

void BallTracker::setOutputPaths(const std::string& video_path, const std::string& csv_path) {
    video_output_path_ = video_path;
    csv_output_path_ = csv_path;
    
    // 确保输出目录存在
    std::string csv_dir = getDirectoryPath(csv_path);
    if (!csv_dir.empty() && !directoryExists(csv_dir)) {
        std::cout << "Creating directory: " << csv_dir << std::endl;
        if (!createDirectory(csv_dir)) {
            std::cerr << "Warning: Failed to create directory: " << csv_dir << std::endl;
        }
    }
    
    // 打开CSV文件
    csv_file_.open(csv_path);
    if (csv_file_.is_open()) {
        csv_file_ << "frame,x,y,confidence\n";
        std::cout << "CSV output opened: " << csv_path << std::endl;
    } else {
        std::cerr << "Failed to open CSV file: " << csv_path << std::endl;
    }
    
    // 初始化视频写入器
    if (!initVideoWriter(video_path, video_width_, video_height_, 25)) { // 使用25fps作为默认
        std::cerr << "Warning: Failed to initialize video writer" << std::endl;
    }
}

bool BallTracker::writeOutputs() {
    if (csv_file_.is_open()) {
        csv_file_.close();
        std::cout << "CSV file closed: " << csv_output_path_ << std::endl;
    }
    
    if (video_writer_.isOpened()) {
        video_writer_.release();
        std::cout << "Video writer released: " << video_output_path_ << std::endl;
    }
    
    return true;
}

void BallTracker::drawTrackingResult(cv::Mat& frame) {
    if (current_position_.confidence > 0.3f) {
        // 绘制球的位置
        cv::circle(frame, cv::Point(current_position_.x, current_position_.y), 8, 
                  cv::Scalar(0, 255, 0), -1);
        cv::circle(frame, cv::Point(current_position_.x, current_position_.y), 8, 
                  cv::Scalar(255, 255, 255), 2);
        
        // 绘制轨迹
        for (size_t i = 1; i < track_history_.size(); i++) {
            if (track_history_[i].confidence > 0.3f && track_history_[i-1].confidence > 0.3f) {
                cv::line(frame, 
                        cv::Point(track_history_[i-1].x, track_history_[i-1].y),
                        cv::Point(track_history_[i].x, track_history_[i].y),
                        cv::Scalar(0, 255, 255), 2);
            }
        }
        
        // 添加信息文本
        std::string info_text = "Frame: " + std::to_string(current_position_.frame_id) + 
                               " Ball: (" + std::to_string(current_position_.x) + ", " + 
                               std::to_string(current_position_.y) + ")";
        cv::putText(frame, info_text, cv::Point(10, 30), 
                   cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(255, 255, 255), 2);
        
        // 添加置信度文本
        std::string conf_text = "Confidence: " + std::to_string(current_position_.confidence).substr(0, 4);
        cv::putText(frame, conf_text, cv::Point(10, 60), 
                   cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(255, 255, 255), 2);
    } else {
        // 没有检测到球
        std::string info_text = "Frame: " + std::to_string(current_position_.frame_id) + " No ball detected";
        cv::putText(frame, info_text, cv::Point(10, 30), 
                   cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 0, 255), 2);
    }
    
    // 添加轨迹点数信息
    std::string track_text = "Track points: " + std::to_string(track_history_.size());
    cv::putText(frame, track_text, cv::Point(10, 90), 
               cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(255, 255, 255), 2);
}
